
--[[
 * rtp_h264_extractor.lua
 * wireshark plugin to extract h264 stream from RTP packets
 * 
 * Copyright (C) 2015 Volvet Zhang <volvet2002@gmail.com>
 *
 * rtp_h264_extractor is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * rtp_h264_extractor is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 *]]


do
    local MAX_JITTER_SIZE = 50
    local rtp_data = Field.new("rtp.payload")
    local rtp_seq = Field.new("rtp.seq")
	
    local function extract_playload_from_rtp()
        local function dump_filter(fd)
            local fh = "rtp";
            if fd ~= nil and fd ~= "" then
                return string.format("%s and (%s)", fh, fd)
            else    
                return fh
            end
        end

        local rtp_tap = Listener.new("ip", dump_filter(get_filter()))
        local text_window = TextWindow.new("rtp playload extractor")
        local filename = ""
        local seq_payload_table = { }
        local pass = 0
        local packet_count = 0
        local max_packet_count = 0
        local fu_info = nil
		
        local function log(info)
            text_window:append(info)
            text_window:append("\n")
        end
        
        -- get_preference is only available since 3.5.0
        if get_preference then
            filename = get_preference("gui.fileopen.dir") .. "/" .. os.date("video_%Y%m%d-%H%M%S.playload")
        else
            filename = "dump.playload"
        end
        
        log("Dumping rtp playload stream to " .. filename)
        local fp = io.open(filename, "wb")
        if fp == nil then 
            log("open dump file fail")
        end
        
        local function on_rtp_payload(seq, payload)
            log("on_rtp_payload:  seq = "..tostring(seq.value)..", payload len = "..tostring(payload.len))
            fp:write(payload.value:tvb()():raw())
            fp:flush()
        end
        
        function rtp_tap.packet(pinfo, tvb)
            local payloadTable = { rtp_data() }
            local seqTable = { rtp_seq() }
            
            if (#payloadTable) < (#seqTable) then 
                log("ERROR: payloadTable size is "..tostring(#payloadTable)..", seqTable size is "..tostring(#seqTable))
                return
            end
            
            if pass == 0 then 
                for i, payload in ipairs(payloadTable) do
                    max_packet_count = max_packet_count + 1
                end
            else 
                packet_count = packet_count + 1
                
                for i, payload in ipairs(payloadTable) do
                    on_rtp_payload(seqTable[1], payload)
                end
                
           
            end 
        end
		
        function rtp_tap.reset()
        end
		
        function rtp_tap.draw() 
        end
		
        local function remove() 
            if fp then 
                fp:close()
                fp = nil
            end
            rtp_tap:remove()
        end 
		
        log("Start")
		
        text_window:set_atclose(remove)
		
        log("phase 1")
        pass = 0
        retap_packets()
        
        log("phase 2:  max_packet_count = "..tostring(max_packet_count))
        pass = 1
        retap_packets()

        if fp ~= nil then 
           fp:close()
           fp = nil
           log("Video stream written to " .. filename)
        end
        
        log("End")
	end


	register_menu("Extract rtp playload stream from RTP", extract_playload_from_rtp, MENU_TOOLS_UNSORTED)
end